import copy
import os.path as osp
import random

import numpy as np
from lib.utils.get_config import get_cfg
from lib.utils.set_seed import seed_everything

cfg = get_cfg(osp.join(osp.abspath("./"), 'configs/config.yaml'))
seed_everything(cfg['SEED'])


class Particle:
    def __init__(self):
        # neural architecture search space
        self.hidden_dim = random.randint(cfg['HIDDEN_DIM_LB'], cfg['HIDDEN_DIM_UB'])
        self.hidden_num = random.randint(cfg['HIDDEN_NUM_LB'], cfg['HIDDEN_NUM_UB'])
        # hyperparameter search space
        self.lr = round(random.uniform(cfg['LEARNING_RATE_LB'], cfg['LEARNING_RATE_UB']), 4)
        # self.parameters = NN_parameters()
        self.cognitiveCoef = 1  # can be changed
        self.socialCoef = 1  # can be changed
        self.informantList = []
        self.informants_best_err = -1
        self.best_err = -1
        self.best_wb = []
        self.informants_best = [self.hidden_num, self.hidden_dim, self.lr]
        self.err = -1  # Current error (set to -1 at start
        self.velocity_hidden_num = random.random()
        self.velocity_hidden_dim = random.random()
        self.velocity_lr = round(random.uniform(0.0100, 0.0001), 4)

    def setInformants(self, swarm, informantNum, index):
        i = 0
        swarm_buffer = copy.deepcopy(swarm)
        banned_index = [index]
        while i < informantNum:
            informant_chosen = np.random.randint(0, len(swarm_buffer))
            if informant_chosen in banned_index:
                continue
            self.informantList.append(swarm[informant_chosen])
            banned_index.append(informant_chosen)
            i += 1

    def set_informant_best(self):
        for informer in self.informantList:
            if informer.best_err < self.informants_best_err or self.informants_best_err == -1:
                self.informants_best_err = informer.best_err
                self.informants_best = informer.best_wb

    def check_error(self, loss):
        self.err = 0
        if self.err < self.best_err or self.best_err == -1:
            self.best_err = self.err
            self.best_wb = [self.hidden_num, self.hidden_dim, self.lr]

    def update_velocity(self):
        inertia_weight = 1

        # Change the velocity values for hidden_num
        r1 = random.random()
        r2 = random.random()

        vel_cog = self.cognitiveCoef * r1 * (self.best_wb[0] - self.hidden_num)
        vel_soc = self.socialCoef * r2 * \
            (self.informants_best[0] - self.hidden_num)
        self.velocity_hidden_num = inertia_weight * \
            self.velocity_hidden_num + vel_soc + vel_cog

        # Change the velocity values for hidden_dim
        r1 = random.random()
        r2 = random.random()

        vel_cog = self.cognitiveCoef * r1 * (self.best_wb[1] - self.hidden_dim)
        vel_soc = self.socialCoef * r2 * \
            (self.informants_best[1] - self.hidden_dim)
        self.velocity_hidden_dim = inertia_weight * \
            self.velocity_hidden_dim + vel_soc + vel_cog

        # Change the velocity values for lr
        r1 = random.random()
        r2 = random.random()

        vel_cog = self.cognitiveCoef * r1 * (self.best_wb[2] - self.lr)
        vel_soc = self.socialCoef * r2 * (self.informants_best[2] - self.lr)
        self.velocity_lr = inertia_weight * self.velocity_lr + vel_soc + vel_cog

    # Update the hidden num and dim
    def change_wb(self):
        self.hidden_num = int(
            round(self.velocity_hidden_num + self.hidden_num))
        self.hidden_dim = int(
            round(self.velocity_hidden_dim + self.hidden_dim))
        self.lr = self.velocity_lr + self.lr
        # keep the particle coordinate within the boundaries
        self.hidden_num = max(self.hidden_num, 2)
        self.hidden_num = min(self.hidden_num, 10)
        self.hidden_dim = max(self.hidden_dim, 2)
        self.hidden_dim = min(self.hidden_dim, 50)
        if self.lr <= 0:
            self.lr = 0.0001
        elif self.lr > 1:
            self.lr = 1.0000
